#!/bin/bash

set -e

source "$(
  cd "$(dirname "$0")" || exit 1
  pwd
)/../common.sh"

#下载依赖和工具包
function install_development() {
  
  source /etc/profile
}

#获取cudnn
function download_cudnn() {
  echo "# cudnn-8.2正在下载"
  CUDNN_TAR_FILE="cudnn-11.4-linux-aarch64sbsa-v8.2.2.26.tgz"
  wget -t 40 -c -P "$1" \
    http://developer.download.nvidia.com/compute/redist/cudnn/v8.2.2/${CUDNN_TAR_FILE} 
  #cp -r  "$2"/src/cudnn-11.4-linux-aarch64sbsa-v8.2.2.26.tgz "$1"/
  if [[ $? -eq 0 ]]; then
    echo -e "\033[1;32;1m# cudnn下载成功\033[0m"
  else
    echo -e "\033[31m# cudnn下载失败,请重新执行脚本\033[0m"
    exit 1
  fi
}

function install_cudnn() {
  echo "# cudnn正在安装"
  tar -zxvf "$1"/cudnn-11.4-linux-aarch64sbsa-v8.2.2.26.tgz -C "$1" && rm -rf "$1"/cudnn-11.4-linux-aarch64sbsa-v8.2.2.26.tgz
  sudo cp "$1"/cuda/include/cudnn*.h /usr/local/cuda-11.4/include 
  sudo cp "$1"/cuda/lib64/libcudnn*  /usr/local/cuda-11.4/lib64
  sudo chmod a+r  /usr/local/cuda-11.4/include/cudnn.h /usr/local/cuda-11.4/lib64/libcudnn*
  echo -e "\033[1;32;1m# cuda-11.4已经安装完毕\033[0m"
}


function main() {
    install_development
    download_cudnn "$1" "$4"
    install_cudnn "$1" 
}

main "$@"